##########################################################################################
## https://sunduanchen.github.io/Scissor/vignettes/Scissor_Tutorial.html
library(Scissor)
library(data.table)
library(optparse)
library(ggplot2)
library(dplyr)
library(patchwork)
library(pROC)
library(parallel)
library(progress)
library(Matrix)

##########################################################################################

option_list <- list(
    make_option(c("--sample_list_file"), type = "character"),
    make_option(c("--use_sample"), type = "character"),
    make_option(c("--use_class"), type = "character"),
    make_option(c("--sample_list_single_file"), type = "character"),
    make_option(c("--sample_list_public_file"), type = "character"),
    make_option(c("--rsem_file"), type = "character"),
    make_option(c("--maf_public_file"), type = "character"),
    make_option(c("--single_cell_file"), type = "character"),
    make_option(c("--gene"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){

    gene <- "TP53"
    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic/"
    sample_list_file <- "~/20220915_gastric_multiple/dna_combinePublic/config/tumor_normal.class.list"
    sample_list_public_file <- paste(work_dir,"/public_ref/combine/MutationInfo.combine.tsv",sep="")
    sample_list_single_file <- "~/20220915_gastric_multiple/dna_combinePublic/config/singleCell_Sample.list"

    rsem_file <- "~/20220915_gastric_multiple/dna_combinePublic/mRNA/CombineTPM.DNAUse.NJMU_TCGA.tsv"
    maf_public_file <- paste(work_dir,"/maf_public/All_use.addVAF.maf",sep="")

    single_cell_file <- paste0(work_dir,"/public_ref/singleCell/njmu/epiall_nor_PCA_50_RE0.5.Rdata")
    out_path <- paste0("~/20220915_gastric_multiple/dna_combinePublic/images/singleCell_MUC6/" , gene)

}

##########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

sample_list_file <- opt$sample_list_file
use_sample <- opt$use_sample
use_class <- opt$use_class
sample_list_single_file <- opt$sample_list_single_file
out_path <- opt$out_path
rsem_file <- opt$rsem_file
sample_list_public_file <- opt$sample_list_public_file
maf_public_file <- opt$maf_public_file
gene <- opt$gene
single_cell_file <- opt$single_cell_file

dir.create(out_path , recursive = T)

##########################################################################################

info <- data.frame(fread(sample_list_file))
info_public <- data.frame(fread(sample_list_public_file))
info_singleCell <- data.frame(fread(sample_list_single_file))

dat_expression <- data.frame(fread(rsem_file))
dat_maf_public <- data.frame(fread( maf_public_file ))
sc_dataset <- load(single_cell_file, verbose = F)
sc_dataset_all <- epiall_nor_PCA_50_RE0.5
##Idents函数定义要取的类别
Idents(sc_dataset_all) <- "sample"   

##########################################################################################

Variant_Types <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

#dat_expression <- subset( dat_expression , gene_id == gene )
dat_maf_public <- subset( dat_maf_public , Hugo_Symbol == gene & Variant_Classification %in% Variant_Types )

colnames(dat_expression) <- gsub( "[.]" , "-" , colnames(dat_expression) )
rownames(dat_expression) <- dat_expression$gene_id
dat_expression <- dat_expression[,-which(colnames(dat_expression)=="gene_id")]

##########################################################################################

info_public <- subset( info_public , From != "NJMU" )
info_public$ID <- info_public$Tumor
info_public$Class_sub <- info_public$Class
info_use <- rbind( info_public[,c( "ID" , "Tumor" , "Class" , "Class_sub" , "From")] , info[,c( "ID" , "Tumor" , "Class" , "Class_sub" , "From")] )

##########################################################################################

dat_maf_public_use <- dat_maf_public

## 突变型样本
mutTumor <- unique(dat_maf_public_use$Tumor_Sample_Barcode)
info_mut <- subset( info_use , Tumor %in% mutTumor )
info_mut <- paste0(info_mut$ID , "_" , info_mut$Class_sub)
info_mut <- info_mut[info_mut %in% colnames(dat_expression)]

## 野生型样本
info_wild <- subset( info_use , !(Tumor %in% mutTumor) )
info_wild <- paste0(info_wild$ID , "_" , info_wild$Class_sub)
info_wild <- info_wild[info_wild %in% colnames(dat_expression)]

##########################################################################################
## 单细胞突变型的样本
info_singleCell_mut <- info_singleCell[info_singleCell$Tumor %in% dat_maf_public_use$Tumor_Sample_Barcode,]
info_singleCell_mut <- subset( info_singleCell_mut , singlecell_ID != "" )

##########################################################################################
## 改成可用多线程
test_logit <- function(X, Y, network, alpha, cell_num, n = 100, nfold = 10){

    set.seed(2)
    m1 <- sum(Y == 1)
    m2 <- sum(Y == 0)
    index1 <- sample(cut(seq(m1), breaks = nfold, labels = F))
    index2 <- sample(cut(seq(m2), breaks = nfold, labels = F))

    print("|**************************************************|")
    print("Perform cross-validation on X with true label")
    AUC_test_real <- c()
    pb1 <- progress_bar$new(total = nfold)

    result <- Reduce(function(x,y)rbind(x,y),mclapply( 1:nfold , function(j){
        c_index <- c(which(Y == 1)[which(index1 == j)], which(Y == 0)[which(index2 == j)])
        X_train <- X[-c_index,]
        Y_train <- Y[-c_index]
        fit <- NULL
        while (is.null(fit$fit)){
            set.seed(123)
            fit <- APML1(X_train, Y_train, family = "binomial", penalty = "Net", alpha = alpha, Omega = network, nlambda = 100)
        }
        index <- which.min(abs(fit$fit$nzero - cell_num))
        Coefs <- as.numeric(fit$Beta[2:(ncol(X_train)+1), index])
        Cell1 <- Coefs[which(Coefs > 0)]
        Cell2 <- Coefs[which(Coefs < 0)]

        X_test <- X[c_index,]
        Y_test <- Y[c_index]
        score_test <- 1/(1+exp(-X_test%*%Coefs-fit$Beta[1,index]))[,1]
        #AUC_test_real[j] <- roc(Y_test, score_test, direction = "<", quiet = T)$auc

        tmp <- data.frame(j = j , auc = roc(Y_test, score_test, direction = "<", quiet = T)$auc)
        tmp

    },mc.cores=1))

    AUC_test_real <- result$auc[order(result$j)]

    cat("Finished!\n")
    print("|**************************************************|")
    print("Perform cross-validation on X with permutated label")
    AUC_test_back <- list()
    pb2 <- progress_bar$new(total = n)

    for (i in 1:n){
        set.seed(i+100)
        AUC_test_back[[i]] <- matrix(0, nfold, 1, dimnames = list(paste0("Testing_",  1:nfold), "AUC"))
        Y2 <- sample(Y)
        names(Y2) <- rownames(X)

        result <- Reduce(function(x,y)rbind(x,y),mclapply( 1:nfold , function(j){
            c_index <- c(which(Y2 == 1)[which(index1 == j)], which(Y2 == 0)[which(index2 == j)])
            X_train <- X[-c_index,]
            Y_train <- Y2[-c_index]
            fit <- NULL
            while (is.null(fit$fit)){
                set.seed(123)
                fit <- APML1(X_train, Y_train, family = "binomial", penalty = "Net", alpha = alpha, Omega = network, nlambda = 100)
            }
            index <- which.min(abs(fit$fit$nzero - cell_num))
            Coefs <- as.numeric(fit$Beta[2:(ncol(X_train)+1), index])
            Cell1 <- Coefs[which(Coefs > 0)]
            Cell2 <- Coefs[which(Coefs < 0)]

            X_test <- X[c_index,]
            Y_test <- Y2[c_index]
            score_test <- 1/(1+exp(-X_test%*%Coefs-fit$Beta[1,index]))[,1]
            tmp <- data.frame(j = j , auc = roc(Y_test, score_test, direction = "<", quiet = T)$auc)
            tmp

        },mc.cores=1))

        AUC_test_back[[i]] <- result$auc[order(result$j)]
    }

    statistic  <- mean(AUC_test_real)
    background <- NULL
    for (i in 1:n){
        background[i] <- mean(AUC_test_back[[i]])
    }
    p <- sum(background > statistic)/n

    print(sprintf("Test statistic = %s", formatC(statistic, format = "f", digits = 3)))
    print(sprintf("Reliability significance test p = %s", formatC(p, format = "f", digits = 3)))

    return(list(statistic = statistic,
                p = p,
                AUC_test_real = AUC_test_real,
                AUC_test_back = AUC_test_back))
}

##########################################################################################
## 按照不同病理类型提取
mut_type <- "All"
for( class in use_class ){

    print(class)

    ## 提取特定病理类型的样本
    sc_dataset <- subset(sc_dataset_all , idents=c(class))
    tmp_singleCell <- subset( info_singleCell_mut , Class == class )  

    if(nrow(tmp_singleCell) > 0){
        Idents(sc_dataset) <- "orig.ident"
        sc_dataset <- subset(sc_dataset , idents=use_sample)

        cell_order <- names(table(sc_dataset$celltype)[order(table(sc_dataset$celltype) , decreasing=T)])

        ## https://www.rdocumentation.org/packages/Seurat/versions/3.0.1/topics/FindNeighbors
        sc_dataset <- FindNeighbors(sc_dataset, dims = 1:10)
        sc_dataset <- RunUMAP(object = sc_dataset, dims = 1:10)

        ## 突变和野生型
        mut_sample <- grep( class , info_mut , value = T)
        wild_sample <- grep( class , info_wild , value = T)
        phenotype <- c(rep(1 , length(mut_sample)) , rep(0,length(wild_sample)))
        names(phenotype) <- c(mut_sample , wild_sample)
        tag <- c('wild-type', 'mutant-type')
        
        ## 使用的转录组数据
        bulk_dataset <- as.matrix(dat_expression[,c(mut_sample , wild_sample)])

        ## scissor
        ## The default value of cutoff is 0.2, i.e., the number of the Scissor selected cells should not exceed 20% of total cells in the single-cell data
        out_name <- paste0( out_path , "/Scissor_STAD_" , gene , "_mutation.",class,".",mut_type,".RData"  )
        if(gene=="TP53"){
                infos4 <- Scissor(bulk_dataset, sc_dataset, phenotype, tag = tag, alpha = seq(0.37,1,0.01) , cutoff = 0.20,
                    family = "binomial", Save_file = out_name)
        }else if(gene=="PIK3CA"){
            infos4 <- Scissor(bulk_dataset, sc_dataset, phenotype, tag = tag, alpha = seq(0.65,1,0.01) , cutoff = 0.20,
                         family = "binomial", Save_file = out_name)
        }else if(gene=="ERBB2"){
            infos4 <- Scissor(bulk_dataset, sc_dataset, phenotype, tag = tag, alpha = seq(0.35,1,0.01) , cutoff = 0.20,
                         family = "binomial", Save_file = out_name)
        }else{
            infos4 <- Scissor(bulk_dataset, sc_dataset, phenotype, tag = tag, alpha = seq(0.20,1,0.01) , cutoff = 0.20,
                         family = "binomial", Save_file = out_name)
        }

        Scissor_select <- rep(0, ncol(sc_dataset))
        names(Scissor_select) <- colnames(sc_dataset)
        Scissor_select[infos4$Scissor_pos] <- 1
        Scissor_select[infos4$Scissor_neg] <- 2
        sc_dataset <- AddMetaData(sc_dataset, metadata = Scissor_select, col.name = "scissor")

        ## Reliability significance test
        numbers <- length(infos4$Scissor_pos) + length(infos4$Scissor_neg)

        if(numbers > 0){
            # load(out_name)
            ## out_name <- paste0( out_path , "/Scissor_STAD_" , gene , "_mutation.",class,".choose.",mut_type,".pdf"  )
            ## To determine whether the inferred phenotype-to-cell associations are reliable, we use the function reliability.test to 
            ## perform a reliability significance test. 
            ## The motivation for the reliability significance test is: 
            ## if the chosen single-cell and bulk data are not suitable for the phenotype-to-cell associations, 
            ## the correlations would be less informative and not well associated with the phenotype labels. 
            ## Thus, the corresponding prediction performance would be poor and not be significantly distinguishable from the randomly permutated labels. 
            ## In this tutorial, we test the identified associations in the above applications as examples to show how to run reliability.test.
            ## 检验的p值
            #result1 <- test_logit(X, Y, network, alpha = infos4$para$alpha , cell_num = numbers, n = 100 , nfold = 10)
            #out_name <- paste0( out_path , "/Scissor_STAD_" , gene , "_mutation.",class,".",mut_type,".tsv"  )
            #write.table( data.frame( p = result1$p) , out_name , row.names = F , quote = F , sep = "\t" )
            
            p1 <- DimPlot(sc_dataset, reduction = 'umap', group.by = 'scissor', cols = c('grey','indianred1','royalblue'), pt.size = 1.2, order = c(2,1))
            #out_name <- paste0( out_path , "/Scissor_STAD_" , gene , "_mutation.",class,".cellType.",mut_type,".pdf"  )
            
            p2 <- DimPlot(sc_dataset, reduction = 'umap', group.by = 'celltype' , pt.size = 1.2, order = c(2,1))

            #out_name <- paste0( out_path , "/Scissor_STAD_" , gene , "_mutation.",class,".sample.",mut_type,".pdf"  )
            #pdf(out_name)
            p3 <- DimPlot(sc_dataset, reduction = 'umap', group.by = 'orig.ident',  pt.size = 1.2, order = c(2,1))
            #dev.off()

            plot <- p1 + p2 + p3
            out_name <- paste0( out_path , "/Scissor_STAD_" , gene , "_mutation.",class,".",mut_type,".pdf"  )
            ggsave(file=out_name,plot=plot,width=18,height=6)

            ## 细胞比例的占比
            ## 占比画图
            tmp_data <- data.frame(table(sc_dataset@meta.data$scissor , sc_dataset@meta.data$celltype))
            colnames(tmp_data) <- c("Scissor_Type" , "Cell_Type" , "Cells")
            tmp_data <- tmp_data %>%
            group_by( Cell_Type ) %>%
            summarize( Scissor_Type = Scissor_Type , Cells = Cells , Cells_Rate = Cells/sum(Cells) )

            tmp_data$Scissor_Type <- as.character(tmp_data$Scissor_Type)
            tmp_data$Scissor_Type <- ifelse( tmp_data$Scissor_Type==0 , "Background" , tmp_data$Scissor_Type )
            tmp_data$Scissor_Type <- ifelse( tmp_data$Scissor_Type==1 , "Scissor+" , tmp_data$Scissor_Type )
            tmp_data$Scissor_Type <- ifelse( tmp_data$Scissor_Type==2 , "Scissor-" , tmp_data$Scissor_Type )
            tmp_data$Cell_Type <- factor( tmp_data$Cell_Type , levels = cell_order , order = T )

            ## 去除主细胞,该细胞数量太少且有MUC6突变定义不明确，以访误导
            #tmp_data <- subset( tmp_data , Cell_Type != "Chief" )
            col <- c('grey','indianred1','royalblue')
            names(col) <- c("Background" , "Scissor+" , "Scissor-" )

            p1 <- ggplot(tmp_data,aes(x=Cell_Type,y=Cells_Rate,fill=factor(Scissor_Type))) +
                geom_bar(stat="identity") +
                ylab("Cell Rate") +
                xlab(NULL) +
                theme_bw() +
                theme(panel.background = element_blank(),#设置背影为白色#清除网格线
                    legend.position ='none',
                    legend.title = element_blank() ,
                    panel.grid.major=element_line(colour=NA),
                    legend.text = element_text(size = 8,color="black",face='bold'),
                    axis.text.x = element_text(size = 10,color="black",face='bold' , angle = 45, hjust = 1),
                    axis.text.y = element_text(size = 8,color="black",face='bold'),
                    axis.title.x = element_text(size = 10,color="black",face='bold'),
                    axis.title.y = element_text(size = 10,color="black",face='bold'),
                    strip.text.x = element_text(size = 15,color="black",face='bold'),
                    axis.line = element_line(size = 0.5))  +
                scale_fill_manual(values=c(col))

            p2 <- ggplot(tmp_data,aes(x=Cell_Type,y=Cells,fill=factor(Scissor_Type))) +
                geom_bar(stat="identity") +
                ylab("Cell Counts") +
                xlab(NULL) +
                theme_bw() +
                theme(panel.background = element_blank(),#设置背影为白色#清除网格线
                    legend.position ='right',
                    legend.title = element_blank() ,
                    panel.grid.major=element_line(colour=NA),
                    legend.text = element_text(size = 8,color="black",face='bold'),
                    axis.text.x = element_text(size = 10,color="black",face='bold' , angle = 45, hjust = 1),
                    axis.text.y = element_text(size = 8,color="black",face='bold'),
                    axis.title.x = element_text(size = 10,color="black",face='bold'),
                    axis.title.y = element_text(size = 10,color="black",face='bold'),
                    strip.text.x = element_text(size = 15,color="black",face='bold'),
                    axis.line = element_line(size = 0.5))  +
                scale_fill_manual(values=c(col))

            plot <- p2 + p1
            out_name <- paste0( out_path , "/Scissor_STAD_" , gene , "_mutation.",class,".CellRate.",mut_type,".pdf"  )
            ggsave(file=out_name,plot=plot,width=9,height=6)

            out_name <- paste0( out_path , "/Scissor_STAD_" , gene , "_mutation.",class,".CellRate.",mut_type,".tsv"  )
            write.table( tmp_data , out_name , row.names = F , quote = F , sep = "\t" )

            out_name <- paste0( out_path , "/Scissor_STAD_" , gene , "_mutation.",class,".CellRate.",mut_type,".RData"  )
            save(sc_dataset , file = out_name)
        }
    }
}
